﻿using System;
using System.Collections.Generic;
using System.IO;
using System.IO.Compression;
using System.Net;
using System.Net.Sockets;
using System.Runtime.Serialization;
using System.Runtime.Serialization.Formatters.Binary;
using System.Security.Cryptography;
using PhiStore.Interfaces;
using PhiStore.Packages;
using PhiStore.Policies;
using PhiStore.ReturnCodes;
using PhiStore.Security;
using PhiStore.Util;

namespace PhiStore
{
    /// <summary>
    /// Client class used for accessing the key-value store known as PhiStore
    /// </summary>
    /// <typeparam name="TValue"></typeparam>
    public class PhiClient<TValue> : IDataStore<TValue>
    {
        private bool initialized = false;

        private bool connecting = false;

        private Prefetcher<TValue> prefetcher = null; //for prefetching

        private PhiCredentials credentials;
        private bool encryptConnection = false;
        private bool compressObjects = true;

        private AesCryptoServiceProvider objectEncryption = null; //for encrypting objects
        private AesCryptoServiceProvider aes = null; //for encrypting communication

        private BinaryFormatter formatter = new BinaryFormatter(); //for serialization

        private int Port = -1; // used for the log to differentiate

        private TcpClient client;
        private Stream inputStream;
        private Stream outputStream;

        private IPEndPoint lastKnownEndpoint;

        #region Simple Functionality
        /// <summary>
        /// Shorthand for ListCreate(new List[value], policy);
        /// Note that the policy for MayAppend and MayPop are set to 'false'
        /// </summary>
        /// <param name="value"></param>
        /// <param name="policy"></param>
        public string Create(TValue value, Policy policy)
        {
            Policy copy = policy.Copy();
            copy.MayAppend = false;
            copy.MayPop = false;

            List<TValue> list = new List<TValue>();
            list.Add(value);
            return ListCreate(list, copy);
        }

        /// <summary>
        /// Shorthand for ListCreate(key, List[value], policy);
        /// Note that the policy for MayAppend and MayPop are set to 'false'
        /// </summary>
        /// <param name="key"></param>
        /// <param name="value"></param>
        /// <param name="policy"></param>
        public void Create(string key, TValue value, Policy policy)
        {
            Policy copy = policy.Copy();
            copy.MayAppend = false;
            copy.MayPop = false;

            List<TValue> list = new List<TValue>();
            list.Add(value);
            ListCreate(key, list, copy);
        }

        /// <summary>
        /// Shorthand for ListRead(key, 0)
        /// </summary>
        /// <param name="key"></param>
        public TValue Read(string key)
        {
            return ListRead(key, 0);
        }

        /// <summary>
        /// Shorthand for ListUpdate(key, value, 0)
        /// </summary>
        /// <param name="key"></param>
        /// <param name="value"></param>
        public void Update(string key, TValue value)
        {
            ListUpdate(key, value, 0);
        }

        /// <summary>
        /// Deletes the entire object including policy etc.
        /// </summary>
        /// <param name="key"></param>
        public void Delete(string key)
        {
            Delete action = new Delete(key);
            ReturnCode code = TryRepeat<ReturnCode>(action);

            if (prefetcher != null)
                prefetcher.RemovePrefetchable(key);

            if (code != ReturnCode.Success)
            {
                throw new PhiException(ReturnCodeTranslation.Translate(code));
            }
        }
        #endregion

        #region List Functionality

        /// <summary>
        /// Creates a new object with the specified values.
        /// </summary>
        /// <param name="values">Serializable values</param>
        /// <param name="policy">Policy dictating behavior and capabilities. This can not be altered</param>
        /// <returns>The generated key</returns>
        public string ListCreate(IList<TValue> values, Policy policy)
        {
            if (!IsSerializable(values))
                throw new ArgumentException("values are not serializable");

            List<byte[]> data = new List<byte[]>();
            foreach (TValue value in values)
            {
                data.Add(StreamUtil.Read(Serialize(value)));
            }

            PhiPackage action = new Create(policy, data);

            ReturnCode code = TryRepeat<ReturnCode>(action);
            string GeneratedKey = ReceiveReply<string>();

            if (code != ReturnCode.Success)
                throw new PhiException(ReturnCodeTranslation.Translate(code));

            if (prefetcher != null)
                prefetcher.AddPrefetchable(GeneratedKey, policy);

            return GeneratedKey;
        }

        /// <summary>
        /// Creates a new object with the specified values.
        /// </summary>
        /// <param name="key">A valid key only containing characters specified in "KeyHelper.ValidCharacters"</param>
        /// <param name="values">Serializable values</param>
        /// <param name="policy">Policy dictating behavior and capabilities. This can not be altered</param>
        public void ListCreate(string key, IList<TValue> values, Policy policy)
        {
            if (!IsSerializable(values))
                throw new ArgumentException("Value is not serializable");

            if (!KeyHelper.KeyIsValid(key))
                throw new ArgumentException("Key is not valid. Valid characters are \"" + KeyHelper.ValidCharacters + "\"");

            List<byte[]> data = new List<byte[]>();
            foreach (TValue value in values)
            {
                data.Add(StreamUtil.Read(Serialize(value)));
            }

            PhiPackage action = new Create(key, policy, data);
            ReturnCode code = TryRepeat<ReturnCode>(action);
            ReceiveReply<string>(); // key read here trash it!

            if (code != ReturnCode.Success)
                throw new PhiException(ReturnCodeTranslation.Translate(code));

            if (prefetcher != null)
                prefetcher.AddPrefetchable(key, policy);
        }

        /// <summary>
        /// Shorthand for ListRead(key, index, 1)[0]
        /// </summary>
        /// <param name="key"></param>
        /// <param name="index"></param>
        /// <returns></returns>
        public TValue ListRead(string key, long index)
        {
            return ListRead(key, index, 1)[0];
        }

        /// <summary>
        /// Read a number of elements within an object
        /// </summary>
        /// <param name="key"></param>
        /// <param name="offset"></param>
        /// <param name="length"></param>
        /// <returns>The elements read</returns>
        public IList<TValue> ListRead(string key, long offset, long length)
        {
            List<TValue> result = new List<TValue>();

            PhiPackage action = new Read(key, offset, length);

            if (prefetcher != null && length == 1)
            {
                try
                {
                    TValue prefetchResult;
                    if (prefetcher.Prefetching(key, offset, out prefetchResult))
                        result.Add(prefetchResult);
                }
                catch
                {
                    Log.Instance.AddLogEntry("PhiClient @ " + Port, "Prefetcher Exception");
                    //prefetch probably failed
                }
            }

            if (result.Count == 0)
            {
                ReadReply reply = TryRepeat<ReadReply>(action);

                if (reply.Code == ReturnCode.Success)
                {
                    List<byte[]> objects = reply.Data<List<byte[]>>();
                    for (int i = 0; i < objects.Count; i++)
                    {
                        result.Add(Deserialize<TValue>(objects[i]));
                    }
                }
                else
                    throw new PhiException(ReturnCodeTranslation.Translate(reply.Code));
            }

            if (prefetcher != null && length == 1)
                prefetcher.Prefetch(key, offset + 1);

            return result;
        }

        /// <summary>
        /// Gets the number of elements for an object
        /// </summary>
        /// <param name="key"></param>
        /// <returns>number of elements</returns>
        public long ListCount(string key)
        {
            ReadReply reply = TryRepeat<ReadReply>(new ListCount(key));
            long result = 0L;
            if (reply.Code == ReturnCode.Success)
            {
                result = reply.Data<long>();
            }
            else
            {
                throw new PhiException(ReturnCodeTranslation.Translate(reply.Code));
            }
            return result;
        }

        /// <summary>
        /// Shorthand for ListUpdate(key, List[value], index);
        /// </summary>
        /// <param name="key"></param>
        /// <param name="value"></param>
        /// <param name="index"></param>
        public void ListUpdate(string key, TValue value, long index)
        {
            List<TValue> values = new List<TValue>();
            values.Add(value);
            ListUpdate(key, values, index);
        }

        /// <summary>
        /// Updates several indices in an object
        /// </summary>
        /// <param name="key"></param>
        /// <param name="values">Serializable values</param>
        /// <param name="offset">Offset to begin updating</param>
        public void ListUpdate(string key, IList<TValue> values, long offset)
        {
            if (!IsSerializable(values))
                throw new ArgumentException("values are not serializable");

            List<byte[]> data = new List<byte[]>();
            foreach (TValue value in values)
            {
                data.Add(StreamUtil.Read(Serialize(value)));
            }
            Update uo = new Update(key, offset, data);

            ReturnCode code = TryRepeat<ReturnCode>(uo);
            if (code != ReturnCode.Success)
            {
                throw new PhiException(ReturnCodeTranslation.Translate(code));
            }
        }

        /// <summary>
        /// Shorthand for ListAppend(key, List[value]);
        /// </summary>
        /// <param name="key"></param>
        /// <param name="value"></param>
        public void ListAppend(string key, TValue value)
        {
            List<TValue> values = new List<TValue>();
            values.Add(value);
            ListAppend(key, values);
        }

        /// <summary>
        /// Append a number of values to an object
        /// </summary>
        /// <param name="key"></param>
        /// <param name="values">Serializable values</param>
        public void ListAppend(string key, IList<TValue> values)
        {
            if (!IsSerializable(values))
                throw new ArgumentException("values are not serializable");

            List<byte[]> data = new List<byte[]>();
            foreach (TValue value in values)
            {
                data.Add(StreamUtil.Read(Serialize(value)));
            }

            ListAppend uo = new ListAppend(key, data);
            ReturnCode code = TryRepeat<ReturnCode>(uo);

            if (code != ReturnCode.Success)
                throw new PhiException(ReturnCodeTranslation.Translate(code));
        }

        /// <summary>
        /// Removes the first element in the object and returns it
        /// </summary>
        /// <param name="key"></param>
        /// <returns>The first element in the object</returns>
        public TValue ListPop(string key)
        {
            ListPop lp = new ListPop(key);
            ReadReply reply = TryRepeat<ReadReply>(lp);

            if (reply.Code == ReturnCode.Success)
            {
                byte[] obj = reply.Data<byte[]>();
                return Deserialize<TValue>(obj);
            }
            else
                throw new PhiException(ReturnCodeTranslation.Translate(reply.Code));
        }

        /// <summary>
        /// Shorthand for ListDelete(key, index, 1)
        /// </summary>
        /// <param name="key"></param>
        /// <param name="index"></param>
        public void ListDelete(string key, long index)
        {
            ListDelete(key, index, 1);
        }

        /// <summary>
        /// Deletes a number of elements in an object
        /// </summary>
        /// <param name="key"></param>
        /// <param name="offset"></param>
        /// <param name="length"></param>
        public void ListDelete(string key, long offset, long length)
        {
            ListDelete ld = new ListDelete(key, offset, length);
            ReturnCode code = TryRepeat<ReturnCode>(ld);

            if (code != ReturnCode.Success)
                throw new PhiException(ReturnCodeTranslation.Translate(code));
        }
        #endregion

        /// <summary>
        /// Used for prefetching
        /// </summary>
        internal Policy GetPolicy(string key)
        {
            ReadReply reply = TryRepeat<ReadReply>(new GetPolicy(key));

            if (reply.Code == ReturnCode.Success)
            {
                return reply.Data<Policy>();
            }
            else
                throw new PhiException(ReturnCodeTranslation.Translate(reply.Code));
        }

        #region Initialization and connection

        /// <summary>
        ///  Constructor for PhiClient with anonymous credentials and prefetching and encryption enabled
        /// </summary>
        public PhiClient()
            : this(true, true)
        { }

        /// <summary>
        /// Constructor for PhiClient with anonymous credentials
        /// </summary>
        /// <param name="encryptConnection">Dictates whether or not communication to/from servers is encrypted</param>
        /// <param name="enablePrefecthing">Enable/disable prefetching</param>
        public PhiClient(bool encryptConnection, bool enablePrefecthing)
            : this(encryptConnection, enablePrefecthing, "", "")
        { }

        /// <summary>
        /// Constructor for PhiClient
        /// </summary>
        /// <param name="encryptConnection">Dictates whether or not communication to/from servers is encrypted</param>
        /// <param name="enablePrefecthing">Enable/disable prefetching</param>
        /// <param name="username">Username used for connecting to PhiServer</param>
        /// <param name="password">Password used for connection to PhiServer</param>
        public PhiClient(bool encryptConnection, bool enablePrefecthing, string username, string password)
        {
            if (enablePrefecthing)
            {
                prefetcher = new Prefetcher<TValue>(encryptConnection, username, password);
            }
            this.encryptConnection = encryptConnection;
            this.credentials = new ClientCredentials(username, password);
        }

        /// <summary>
        /// Initializes the PhiClient.
        /// No compression and no object encryption.
        /// </summary>
        public void Initialize()
        {
            Initialize(false, "");
        }

        /// <summary>
        /// Initializes the PhiClient with no object encryption
        /// </summary>
        /// <param name="compress">Dictates whether or not objects should be compressed to preserve storage</param>
        public void Initialize(bool compress)
        {
            Initialize(compress, "");
        }

        /// <summary>
        /// Initializes the PhiClient with no compression
        /// </summary>
        /// <param name="encryptionKey">Key used for encrypting and decrypting objects. ""=no encryption</param>
        public void Initialize(string encryptionKey)
        {
            Initialize(false, encryptionKey);
        }

        /// <summary>
        /// Initializes the PhiClient.
        /// </summary>
        /// <param name="compress">Dictates whether or not objects should be compressed to preserve storage</param>
        /// <param name="encryptionKey">Key used for encrypting and decrypting objects. ""=no encryption</param>
        public void Initialize(bool compress, string encryptionKey)
        {
            Log.Instance.AddLogEntry("PhiClient", "Initializing...");
            initialized = true;

            if (encryptionKey != "")
            {
                this.objectEncryption = new AesCryptoServiceProvider();
                this.objectEncryption.Key = Helper.GenerateKey(encryptionKey, 256);
                this.objectEncryption.IV = Helper.GenerateIV(encryptionKey, 128); //apparently this has to be the same as well
            }
            else
            {
                this.objectEncryption = null;
            }

            if (prefetcher != null)
            {
                prefetcher.Initialize(compress, encryptionKey);
            }
            this.compressObjects = compress;
        }

        #region connection helper methods
        private ReturnCode SendConnectionRequest(RSAParameters key, byte[] connectionKey, byte[] iv)
        {
            RSACryptoServiceProvider rsa = new RSACryptoServiceProvider();
            rsa.ImportParameters(key);

            Connect action = (encryptConnection ?
                new Connect(credentials, connectionKey, iv) :
                new Connect(credentials));

            RSAEncrypted encrypted = Encryption.RSAEncrypt(action, rsa, formatter);

            return TryRepeat<ReturnCode>(encrypted);
        }

        private void initConnect(IPAddress ip, int port)
        {
            connecting = true;

            client = new TcpClient();

            lastKnownEndpoint = new IPEndPoint(ip, port);

            client.Connect(lastKnownEndpoint);
            client.SendTimeout = 30000;
            client.ReceiveTimeout = 30000;

            inputStream = client.GetStream();
            outputStream = client.GetStream();

        }

        private void doConnect(BroadcastObject bc)
        {
            ReturnCode result = ReturnCode.Failure;
            if (encryptConnection)
            {
                aes = new AesCryptoServiceProvider();

                aes.Key = Helper.GenerateKey(256);
                aes.GenerateIV();

                result = SendConnectionRequest(bc.PublicKey, aes.Key, aes.IV);
            }
            else
            {
                result = SendConnectionRequest(bc.PublicKey, null, null);
            }

            connecting = false;
            if (result != ReturnCode.Success)
            {
                throw new WebException(ReturnCodeTranslation.Translate(result));
            }
            else
            {
                int oldPort = Port;
                Port = (client.Client.LocalEndPoint as IPEndPoint).Port;
                Log.Instance.AddLogEntry("PhiClient @ " + Port, "Connection Success (previous port: " + oldPort + ")");
            }
        }
        #endregion

        /// <summary>
        /// Connects directly to a remote server
        /// </summary>
        /// <param name="ip">IPAddress of the server</param>
        /// <param name="port">Port on which to connect</param>
        public void Connect(IPAddress ip, int port)
        {
            if (!initialized)
            {
                throw new ArgumentException("Store not initialized. Call Initialize first!");
            }

            Log.Instance.AddLogEntry("PhiClient", "Connecting directly to IP:" + ip.ToString() + ":" + port);

            try
            {
                initConnect(ip, port);

                SendRequest(new WhoAreYou());
                BroadcastObject bc = ReceiveReply<BroadcastObject>();

                doConnect(bc);

                if (prefetcher != null)
                {
                    prefetcher.Connect(ip, port);
                }

            }
            catch (Exception e)
            {
                throw new WebException("Connection failed", e);
            }

        }

        /// <summary>
        /// Tries to receive a multicast @ port, and connect afterwards
        /// </summary>
        /// <param name="multicastPort">port specified</param>
        public void Connect(int multicastPort)
        {
            Log.Instance.AddLogEntry("PhiClient", "Connecting by UDP multicast");
            if (!initialized)
            {
                throw new ArgumentException("Store not initialized. Call Initialize first!");
            }

            try
            {
                BroadcastObject bc = ReceiveBroadcast.Receive(multicastPort, 5);
                initConnect(bc.IP, bc.Port);

                doConnect(bc);

                if (prefetcher != null)
                {
                    prefetcher.Connect(bc.IP, bc.Port);
                }
            }
            catch (Exception e)
            {
                throw new WebException("Connection failed", e);
            }

        }

        /// <summary>
        /// Shorthand for Connect(24281)
        /// </summary>
        public void Connect()
        {
            Connect(24281);
        }

        /// <summary>
        /// Shuts down the connection
        /// Shuts down the prefetcher
        /// </summary>
        public void Disconnect()
        {
            if (client != null && client.Connected)
            {
                try
                {
                    SendRequest(new Disconnect());
                }
                catch { }
                client.Close();
                client = null;
            }

            if (prefetcher != null)
            {
                prefetcher.Stop();
                prefetcher = null;
            }
        }

        /// <summary>
        /// Disconnects
        /// </summary>
        ~PhiClient()
        {

            Disconnect();
        }

        #endregion

        #region helper methods
        private T TryRepeat<T>(PhiPackage action)
        {
        REDO:
            object result = ReturnCode.WorkingOnObject;
            try
            {
                if (!initialized)
                {
                    throw new ArgumentException("Store not initialized. Call Initialize first!");
                }
                else if (client == null)
                {
                    throw new ArgumentException("Store not connected. Call connect first!");
                }

                SendRequest(action);
            }
            catch (Exception e)
            {
                if (e is WebException || e is IOException)
                {
                    if (!(action is Disconnect) && Reconnect()) //reconnect success;
                    {
                        Log.Instance.AddLogEntry("PhiClient @ " + Port, "Reconnect Success");
                        goto REDO;
                    }
                }
                else
                    throw e;
            }
            try
            {   //receive reply
                while (ReturnCode.WorkingOnObject.Equals(result))
                {
                    result = ReceiveReply<object>();

                    if (ReturnCode.WorkingOnObject.Equals(result))
                        Log.Instance.Beep();
                }
            }
            catch (Exception e)
            {
                if (e is WebException || e is IOException)
                {
                    if ((action is Create) ||
                        (action is Delete) ||
                        (action is ListAppend) ||
                        (action is ListDelete) ||
                        (action is ListPop) ||
                        (action is Update))
                    {
                        throw new PhiException("A modifying operation failed - Unable to know if it was performed or not", e);
                    }
                    else if (!(action is Disconnect) && Reconnect()) //reconnect success;
                    {
                        Log.Instance.AddLogEntry("PhiClient @ " + Port, "Reconnect Success");
                        goto REDO;
                    }
                }
                else
                    throw e;
            }
            return (T)result;
        }

        private void SendRequest(PhiPackage action)
        {
            Log.Instance.AddLogEntry("PhiClient @ " + Port, "Sending " + action.ToString() + " to " + client.Client.RemoteEndPoint.ToString());
            PhiPackage toSend = action;
            if (aes != null && encryptConnection && !(action is RSAEncrypted))
            {
                toSend = Encryption.Encrypt(toSend, aes, formatter);
            }

            formatter.Serialize(outputStream, toSend);
            outputStream.Flush();
        }

        private T ReceiveReply<T>()
        {
            object reply = new object();
            try
            {
                reply = formatter.Deserialize(inputStream);
                if (reply is Encrypted)
                    reply = Encryption.Decrypt((Encrypted)reply, aes, formatter);
            }
            catch (IOException e)
            {
                throw new WebException("Connection was lost ", e);
            }
            catch (SerializationException e)
            {
                throw new WebException("Connection was lost ", e);
            }

            return (T)reply;
        }

        private bool Reconnect()
        {
            bool success = false;
            if (!connecting) // we are already trying to connect
            {
                aes = null;
                Log.Instance.AddLogEntry("PhiClient @ " + Port, "Reconnecting");
                try
                {
                    Connect(lastKnownEndpoint.Address, lastKnownEndpoint.Port);
                    success = true;
                }
                catch (WebException e) //direct connect failed
                {
                    try
                    {
                        Connect(); //trying multicast
                        success = true;
                    }
                    catch (WebException ex)
                    {
                        if (prefetcher != null)
                            prefetcher.Stop();

                        throw new WebException("Reconnection failed: " + e.ToString() + " [AND] " + ex.ToString());
                    }
                }
            }
            return success;
        }

        /// <summary>
        /// Dirty hack for compression: http://www.damirscorner.com/AlwaysCloseDeflateStreamBeforeReadingResults.aspx
        /// </summary>
        /// <param name="o">object to serialize</param>
        /// <returns>a stream containing the serialized object</returns>
        private Stream Serialize(object o)
        {
            MemoryStream stream = new MemoryStream();

            if (compressObjects && objectEncryption != null)
            { //compress and encrypt
                CryptoStream crypt = new CryptoStream(stream, objectEncryption.CreateEncryptor(), CryptoStreamMode.Write);
                DeflateStream compressor = new DeflateStream(crypt, CompressionMode.Compress, true);
                formatter.Serialize(compressor, o);
                compressor.Flush();
                compressor.Close(); //HACK
                crypt.FlushFinalBlock();
            }
            else if (compressObjects)
            {//compress only
                DeflateStream compressor = new DeflateStream(stream, CompressionMode.Compress, true);
                formatter.Serialize(compressor, o);
                compressor.Flush();
                compressor.Close(); //HACK
            }
            else if (objectEncryption != null)
            {
                CryptoStream crypt = new CryptoStream(stream, objectEncryption.CreateEncryptor(), CryptoStreamMode.Write);
                formatter.Serialize(crypt, o);
                crypt.FlushFinalBlock();
            }
            else
            {
                formatter.Serialize(stream, o);
                stream.Flush();
            }
            stream.Position = 0;

            return stream;
        }

        private T Deserialize<T>(byte[] data)
        {
            Stream stream = new MemoryStream();
            StreamUtil.Copy(data, stream);
            stream.Position = 0;

            if (objectEncryption != null)
            {
                stream = new CryptoStream(stream, objectEncryption.CreateDecryptor(), CryptoStreamMode.Read);
            }
            if (compressObjects)
            {
                stream = new DeflateStream(stream, CompressionMode.Decompress);
            }

            T result = (T)formatter.Deserialize(stream);

            stream = null;
            return result;
        }

        private bool IsSerializable(IEnumerable<object> os)
        {
            bool result = true;
            foreach (object o in os)
            {
                result &= IsSerializable(o);
            }
            return result;
        }

        private bool IsSerializable(object o)
        {
            return o.GetType().IsSerializable;
        }

        #endregion
    }
}
